//
// Created by benxb on 2021/12/12.
//

#include <index/linear_index_segment.h>
#include <common/exception.h>
#include <utility>
#include <common/util/file_util.h>

namespace benxdb {

LinearIndexSegment::~LinearIndexSegment() {
    index_io_.close();
}

LinearIndexSegment::LinearIndexSegment(const std::string index_dir, index_seg_id_t seg_id, LinearIndexKeyType type) {
    key_type_ = type;
    index_dir_ = index_dir;
    index_file_ = index_dir + "/" + std::to_string(seg_id);
    InitHeader(seg_id, type);
    ReadSegmentFile();
}

void LinearIndexSegment::ReaderHeader() {
    index_io_.seekp(0);
    index_io_.read(reinterpret_cast<char *>(&header_), sizeof(header_));
}

void LinearIndexSegment::SetSegmentId(index_seg_id_t seg_id) {
    header_.seg_id = seg_id;
    index_io_.seekp(OFFSET_SEGMENT_ID);
    index_io_.write(reinterpret_cast<const char *>(&seg_id), sizeof(seg_id));
    index_io_.flush();
}

index_seg_id_t LinearIndexSegment::GetSegmentId() const {
    return header_.seg_id;
}

void LinearIndexSegment::SetNextSegmentId(index_seg_id_t next_seg) {
    header_.next_seg_id = next_seg;
    index_io_.seekp(OFFSET_NEXT_SEGMENT_ID);
    index_io_.write(reinterpret_cast<const char *>(&next_seg), sizeof(next_seg));
    index_io_.flush();
}

index_seg_id_t LinearIndexSegment::GetNextSegmentId() const {
    return header_.next_seg_id;
}

void LinearIndexSegment::WriteHeader() {
    index_io_.seekp(0);
    index_io_.write(reinterpret_cast<const char *>(&header_), sizeof(header_));
    index_io_.flush();
}

void LinearIndexSegment::SetPageCount(uint32_t page_count) {
    header_.page_count = page_count;
    WriteHeader();
}

uint32_t LinearIndexSegment::GetPageCount() const {
    return header_.page_count;
}

void LinearIndexSegment::AppendIndexPair(LinearIndexLeafPair &pair) {
    index_io_.seekp(0, std::ios_base::end);
    index_io_.write(pair.key.GetVal(), pair.key.GetSize());
    index_io_.write(reinterpret_cast<const char *>(&pair.record_id), sizeof(pair.record_id));
    index_io_.flush();

    if (pair.key.CompareLessThan(header_.min_key)) {
        header_.min_key = pair.key;
    }
    if (pair.key.CompareGreaterThan(header_.max_key)) {
        header_.max_key = pair.key;
    }
    index_file_size_ += pair.GetSize();
    header_.tmp_page.size += pair.GetSize();
    if (header_.tmp_page.size + pair.GetSize() > PAGE_SIZE - LinearIndexPage::GetHeaderSize()) {
        CompressTmpPage();
    }
    WriteHeader();
}

void LinearIndexSegment::CompressTmpPage() {
    index_file_size_ -= header_.tmp_page.size;
    int page_size = header_.tmp_page.size;
    auto* page = new LinearIndexPage(key_type_);
    index_io_.seekp(header_.tmp_page.offset);
    index_io_.read(page->GetData() + LinearIndexPage::GetHeaderSize(), page_size);
    page->SetIndexSize(page_size/LinearIndexLeafPair::SizeOfKeyType(key_type_));
    page->SetKeyType(key_type_);
    page->SetPageID(0);
    // TODO compress
    int compressed_size = page_size + LinearIndexPage::GetHeaderSize();
    char* compressed_data = page->GetData();
    // append index page
    index_io_.seekp(header_.tmp_page.offset);
    index_io_.write(compressed_data, compressed_size);

    header_.page_ids[header_.page_count].offset = header_.tmp_page.offset;
    header_.page_ids[header_.page_count].size = compressed_size;
    header_.page_count++;
    index_file_size_ += compressed_size;
    header_.tmp_page.offset = index_file_size_;
    header_.tmp_page.size = 0;
}

bool LinearIndexSegment::GetIndexPage(uint32_t page_index, LinearIndexPage &page) {
    LinearIndexOffsetSize off_size{};
    if (page_index < header_.page_count) {
        off_size = header_.page_ids[page_index];
        index_io_.seekp(off_size.offset);
        index_io_.read(page.GetData(), off_size.size);
    } else if (page_index == header_.page_count && header_.tmp_page.size > 0) {
        // read tmp page
        off_size = header_.tmp_page;
        index_io_.seekp(off_size.offset);
        index_io_.read(page.GetData() + benxdb::LinearIndexPage::GetHeaderSize(), off_size.size);
        page.SetIndexSize(off_size.size/LinearIndexLeafPair::SizeOfKeyType(key_type_));
        page.SetKeyType(key_type_);
        page.SetPageID(0);
    } else {
        return false;
    }
    return true;
}

void LinearIndexSegment::AppendManyIndexPairs(std::vector<LinearIndexLeafPair> &pairs) {
    if (pairs.empty()) {
        return;
    }
    index_io_.seekp(0, std::ios_base::end);
    for (auto& pair : pairs) {
        AppendSingleIndexWithoutFlush(pair);
    }
    index_io_.flush();
    WriteHeader();
}

void LinearIndexSegment::AppendManyIndexPairs(std::vector<LinearIndexLeafPair> &pairs, int cposL, int cposH) {
    if (pairs.empty()) {
        return;
    }
    index_io_.seekp(0, std::ios_base::end);
    for (int i = cposL; i <= cposH; i++) {
        auto& pair = pairs[i];
        AppendSingleIndexWithoutFlush(pair);
    }
    index_io_.flush();
    WriteHeader();
}

void LinearIndexSegment::AppendSingleIndexWithoutFlush(LinearIndexLeafPair &pair) {
    index_io_.write(pair.key.GetVal(), pair.key.GetSize());
    index_io_.write(reinterpret_cast<const char *>(&pair.record_id), sizeof(pair.record_id));
    if (pair.key.CompareLessThan(header_.min_key)) {
        header_.min_key = pair.key;
    }
    if (pair.key.CompareGreaterThan(header_.max_key)) {
        header_.max_key = pair.key;
    }
    index_file_size_ += pair.GetSize();
    header_.tmp_page.size += pair.GetSize();
    if (header_.tmp_page.size + pair.GetSize() > PAGE_SIZE - LinearIndexPage::GetHeaderSize()) {
        CompressTmpPage();
    }
}

LinearIndexKey LinearIndexSegment::GetApproximateMid()  {
    return header_.min_key.Add(header_.max_key).Divide(LinearIndexKey(LinearIndexKeyType::STYPE_FLOAT, 2.0));
}

LinearIndexSegment::LinearIndexSegment(const std::string tmp_index_file, LinearIndexKeyType type) :
        index_file_(tmp_index_file), key_type_(type) {
    auto dir_end = index_file_.rfind('/');
    index_dir_ = index_file_.substr(0, dir_end);
    InitHeader(0, type);
    ReadSegmentFile();
}

std::shared_ptr<LinearIndexSegment> LinearIndexSegment::CreateTempSegment(int tmp_index) {
    std::string tmp_file_name = index_file_ + ".tmp" + std::to_string(tmp_index);
    return std::make_shared<LinearIndexSegment>(tmp_file_name, key_type_);
}

void LinearIndexSegment::Flush() {
    index_io_.flush();
    WriteHeader();
}

void LinearIndexSegment::Destroy() {
    index_io_.close();
    remove(index_file_.c_str());
}

void LinearIndexSegment::RenameSegment(index_seg_id_t seg_id) {
    std::string new_index_file = index_dir_ + "/" + std::to_string(seg_id);
    index_io_.close();
    rename(index_file_.c_str(), new_index_file.c_str());
    index_file_ = new_index_file;
    index_io_.open(index_file_, std::ios::binary | std::ios::in | std::ios::out);
    if (!index_io_.is_open()) {
        throw Exception("can't open index segment file : " + new_index_file);
    }
    ReaderHeader();
    header_.seg_id = seg_id;
    WriteHeader();
    index_file_size_ = FileUtil::GetFileSize(index_file_);
}

void LinearIndexSegment::ReadSegmentFile() {
    bool new_segment_file = false;
    index_io_.open(index_file_, std::ios::binary | std::ios::in | std::ios::out);
    // directory or file does not exist
    if (!index_io_.is_open()) {
        new_segment_file = true;
        index_io_.clear();
        // create a new file
        index_io_.open(index_file_, std::ios::binary | std::ios::trunc | std::ios::out);
        index_io_.close();
        // reopen with original mode
        index_io_.open(index_file_, std::ios::binary | std::ios::in | std::ios::out);
        if (!index_io_.is_open()) {
            throw Exception("can't open index segment file: " + index_file_);
        }
    }
    if (new_segment_file) {
        // new segment file
        WriteHeader();
        index_file_size_ = sizeof(LinearIndexSegmentHeader);
    } else {
        ReaderHeader();
        index_file_size_ = 0;
    }
}

void LinearIndexSegment::WriteTmpPageHeader() {
    page_id_t page_id = header_.page_count;
    uint32_t index_count = 0;
    index_io_.write(reinterpret_cast<const char *>(&page_id), sizeof(page_id));
    index_io_.write(reinterpret_cast<const char *>(&header_.key_type), sizeof(header_.key_type));
    index_io_.write(reinterpret_cast<const char *>(&index_count), sizeof(index_count));
    index_io_.flush();
}

void LinearIndexSegment::InitHeader(index_seg_id_t seg_id, LinearIndexKeyType type) {
    header_.seg_id = seg_id;
    header_.key_type = type;
    header_.tmp_page.offset = sizeof(LinearIndexSegmentHeader);
    header_.tmp_page.size = 0;
    header_.page_count = 0;
    header_.min_key = LinearIndexKey::GetMaxValueOfType(type);
    header_.max_key = LinearIndexKey::GetMinValueOfType(type);
}

}